import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression

life_sat = pd.read_csv("./datasets/lifesat/oecd_bli_2015.csv", thousands=',')
life_sat_total = life_sat[life_sat["INEQUALITY"] == "TOT"]
life_sat_total = life_sat_total.pivot(index = "Country", columns = "Indicator", values = "Value")

gdp_per_capita = pd.read_csv("./datasets/lifesat/gdp_per_capita.csv", thousands=',', delimiter = '\t', encoding = 'latin1', na_values= "n/a", index_col = "Country", usecols = ['Country', '2015'])
gdp_per_capita.rename(columns = {"2015" : "GDP per capita"}, inplace = True)

full_country_stats = pd.merge(left = life_sat_total, right = gdp_per_capita, left_index = True, right_index = True)
full_country_stats.sort_values(by = "GDP per capita", inplace = True)

remove_indices = [0, 1, 6, 8, 33, 34, 35]
keep_indices = list(set(range(36)) - set(remove_indices))
sample_data = full_country_stats[["GDP per capita", "Life satisfaction"]].iloc[keep_indices]
sample_data.to_csv("./datasets/lifesat/out.csv")

sample_data.plot(kind = "scatter", grid = True, x = "GDP per capita", y = "Life satisfaction")
plt.show()

x = np.c_[sample_data["GDP per capita"]]
y = np.c_[sample_data["Life satisfaction"]]

model = LinearRegression()
model.fit(x, y)

print(model.predict([[22587], [151589]]))  #5.96242338